package tree

// https://leetcode-cn.com/leetbook/read/top-interview-questions-medium/xvuyv3/
// 获得搜索树的位置
// 算法: 分别遍历左右子树的大小，并做计算
func kthSmallest(root *TreeNode, k int) int {
	treeSizeMap := make(map[*TreeNode]int)
	recurForTreeSize(root, treeSizeMap)
	return recurForKthNum(root, k, treeSizeMap).Val
}

func recurForTreeSize(root *TreeNode, treeSizeMap map[*TreeNode]int) int {
	if nil == root {
		return 0
	}
	currentSize := 1 + recurForTreeSize(root.Left, treeSizeMap) + recurForTreeSize(root.Right, treeSizeMap)
	treeSizeMap[root] = currentSize
	return currentSize
}

func recurForKthNum(root *TreeNode, k int, treeSizeMap map[*TreeNode]int) *TreeNode {
	if nil == root {
		return nil
	}
	// 如果传入的key 为nil，返回的也是0
	if 1+treeSizeMap[root.Left] == k {
		return root
	} else if treeSizeMap[root.Left] >= k {
		return recurForKthNum(root.Left, k, treeSizeMap)
	} else {
		return recurForKthNum(root.Right, k-1-treeSizeMap[root.Left], treeSizeMap)
	}
}
